#!/usr/bin/env python
import numpy as np
PI = np.pi


class Gauss1DFeatureExtractor():

    def __init__(self, n_gaussian, n_level=1, sigma=5e-1, mode=1, a=0, b=1):
        """
           mode:控制每层gauss数目
                default(1): 随着层数增加 gauss数目翻倍
                0: 随着层数增加 gauss数目不变
        """
        self.n_gaussian = n_gaussian
        self.n_level = n_level
        self.sigma = sigma
        self.mode = mode
        self.a = a
        self.b = b
        self.info = None  # List[dic]   level: information   level 仍从1开始编号
        # seed
        np.random.seed(1234)
        self.setInfo()

    def setInfo(self):
        self.info = [{} for _ in range(self.n_level + 1)]

        for level in range(1, self.n_level + 1):
            points_num = self.getLevelNum(level)
            sigma = self.getLevelSigma(level)
            self.info[level]["points_num"] = points_num
            self.info[level]["sigma"] = sigma # 随机选择 np.random.uniform(0, 1, points_num)
            # self.info[level]["coordinates"] = (self.a + np.random.uniform(0, 1, points_num) * (self.b - self.a))
            extend = (self.b - self.a) * 0.1
            self.info[level]["coordinates"] = np.linspace(self.a - extend, self.b + extend, points_num)

    def getLevelNum(self, level):
        """
            mode=1, double when level increase with half sigma
        """
        if self.mode == 1:
            return self.n_gaussian * 2 ** (level - 1)
        elif self.mode == 0:
            return self.n_gaussian

    def getLevelSigma(self, level):
        """
            half sigma when level increase
        """
        return self.sigma / 2 ** (level - 1)

    def getSumGaussians(self):
        """
            sum of gaussians
        """
        return sum(group["points_num"] for group in self.info if group)
    
    def deleteLevel(self, level):
        '''
            delete level   points num = 0
        '''
        print(f"Delete level {level}")
        points_num = 0
        self.info[level]["points_num"] = points_num
        extend = (self.b - self.a) * 0.1
        self.info[level]["coordinates"] = np.linspace(self.a - extend, self.b + extend, points_num)

    def updateLevel(self, level, indices):
        '''
             delete some gaussians in (level) through (indices)
        '''
        if self.info[level]["points_num"] == 0:
            print(f"Pass level {level}")
        before_num, points_num = self.info[level]["points_num"], indices.sum()
        print(f"level {level} from {before_num} to {points_num}")
        self.info[level]["points_num"] = points_num
        gaussians = self.info[level]["coordinates"][indices]
        self.info[level]["coordinates"] = gaussians

    def __repr__(self):
        text_init = f'Initial Gaussians: {self.n_gaussian}, initial sigma: {self.sigma}'
        text_gauss = f'Points_num: {[group["points_num"] for group in self.info if group]}, Sum: {self.getSumGaussians()}'
        text_sigma = f'Sigma: {[group["sigma"] for group in self.info if group]}'
        return "\n".join([text_init, text_gauss, text_sigma])

    def eval(self, x_in, eval_list=('u',)):
        if x_in.ndim == 1:
            x_in = x_in[:, np.newaxis]

        Z = []
        Z1 = []
        Z2 = []
        for level in range(1, self.n_level + 1):
            sigma = self.info[level]["sigma"]
            mu = self.info[level]["coordinates"]
            z = np.exp(-0.5 * (x_in - mu) ** 2 / sigma ** 2)
            z1 = z * (-(x_in - mu) / sigma ** 2)
            z2 = z * (-1 / sigma ** 2) + z1 * (-(x_in - mu) / sigma ** 2)
            Z.append(z)
            Z1.append(z1)
            Z2.append(z2)
        Z = np.hstack(Z)
        Z1 = np.hstack(Z1)
        Z2 = np.hstack(Z2)

        eval_all = {}
        for eval_item in eval_list:
            item_order = len(eval_item)
            if item_order == 1:
                eval_all[eval_item] = Z
            elif item_order == 2:
                eval_all[eval_item] = Z1
            elif item_order == 3:
                eval_all[eval_item] = Z2
        return eval_all


class Gauss2DFeatureExtractor():
    def __init__(self, n_gaussian=50, n_level=1, sigma=0.5, mode=1, domain=(-1, 1, -1, 1)):
        """
           mode:控制每层gauss数目
                default(1): 随着层数
                0:
        """
        self.n_gaussian = n_gaussian             # 初始Gauss
        self.n_level = n_level
        self.sigma = sigma                       # 初始sigma
        self.mode = mode
        self.x_min, self.x_max, self.y_min, self.y_max = domain
        self.info = None  # List[dic]   level: information   level 仍从1开始编号
        self.setInfo()
        #####  random seed ####

    def setInfo(self):
        self.info = [{} for _ in range(self.n_level + 1)]

        for level in range(1, self.n_level + 1):
            points_num = self.getLevelNum(level)
            sigma = self.getLevelSigma(level)
            self.info[level]["points_num"] = points_num
            self.info[level]["sigma"] = sigma
            extend_x, extend_y = (self.x_max - self.x_min) * 0.1, (self.y_max - self.y_min) * 0.1 
            x_all = np.random.uniform(self.x_min - extend_x, self.x_max + extend_x, points_num)
            y_all = np.random.uniform(self.y_min - extend_y, self.y_max + extend_y, points_num)
            # x_all = (self.x_min + np.random.uniform(0, 1, points_num) * (self.x_max - self.x_min))
            # y_all = (self.y_min + np.random.uniform(0, 1, points_num) * (self.y_max - self.y_min))
            self.info[level]["coordinates"] = np.hstack([x_all[:, np.newaxis], y_all[:, np.newaxis]])

    def getLevelNum(self, level):
        """
            mode=1, double when level increase with half sigma
        """
        if self.mode == 1:
            return self.n_gaussian * 2 ** (level - 1)
        elif self.mode == 0:
            return self.n_gaussian

    def getLevelSigma(self, level):
        """
            half sigma when level increase
        """
        return self.sigma / 2 ** (level - 1)

    def getSumGaussians(self):
        """
            sum of gaussians
        """
        return sum(group["points_num"] for group in self.info if group)

    def __repr__(self):
        text_init = f'Initial Gaussians: {self.n_gaussian}, initial sigma: {self.sigma}'
        text_gauss = f'Points_num: {[group["points_num"] for group in self.info if group]}, Sum: {self.getSumGaussians()}'
        text_sigma = f'Sigma: {[group["sigma"] for group in self.info if group]}'
        return "\n".join([text_init, text_gauss, text_sigma])

    def eval(self, x_in, eval_list=('u',)):
        """
            A simple version
        """
        if x_in.ndim == 1:
            x_in = x_in[:, np.newaxis]

        x_pde = x_in[:, [0]]
        y_pde = x_in[:, [1]]

        G = []
        Gx = []
        Gy = []
        Gxx = []
        Gyy = []
        Gxy = []

        for level in range(1, self.n_level + 1):
            sigma = self.info[level]["sigma"]
            xy_gauss = self.info[level]["coordinates"]
            x_gauss = xy_gauss[:, [0]].transpose()
            y_gauss = xy_gauss[:, [1]].transpose()
            g = np.exp(- 0.5 / sigma**2 * ((x_pde - x_gauss)**2 + (y_pde - y_gauss)**2))
            gx = - 1 / sigma**2 * (x_pde - x_gauss) * g
            gy = - 1 / sigma**2 * (y_pde - y_gauss) * g
            gxx = - 1 / sigma**2 * (1 - (x_pde - x_gauss)**2 / sigma**2) * g
            gyy = - 1 / sigma**2 * (1 - (y_pde - y_gauss)**2 / sigma**2) * g
            gxy = - 1 / sigma**4 * (x_pde - x_gauss) * (y_pde - y_gauss) * g

            G.append(g)
            Gx.append(gx)
            Gy.append(gy)
            Gxx.append(gxx)
            Gyy.append(gyy)
            Gxy.append(gxy)

        G = np.hstack(G)
        Gx = np.hstack(Gx)
        Gy = np.hstack(Gy)
        Gxx = np.hstack(Gxx)
        Gxy = np.hstack(Gxy)
        Gyy = np.hstack(Gyy)

        eval_all = {}
        for eval_item in eval_list:
            if eval_item == "u":
                eval_all[eval_item] = G
            elif eval_item == "u0":
                eval_all[eval_item] = Gx
            elif eval_item == "u1":
                eval_all[eval_item] = Gy
            elif eval_item == "u00":
                eval_all[eval_item] = Gxx
            elif eval_item == "u01":
                eval_all[eval_item] = Gxy
            elif eval_item == "u11":
                eval_all[eval_item] = Gyy

        return eval_all


if __name__ == "__main__":
    extractor = Gauss2DFeatureExtractor(n_level=3)
    print(extractor)
    xy = np.random.random((20, 2))
    extractor.eval(xy)







